import numpy as np
from tqdm import tqdm
import networkx as nx
import scipy.sparse as sp
from collections import defaultdict

import logging
from collections import defaultdict
import warnings
warnings.filterwarnings('ignore')

n_users = 0
n_items = 0
n_entities = 0
n_relations = 0
n_nodes = 0
train_user_set = defaultdict(list)
test_user_set = defaultdict(list)

# 读入user访问的aor
def read_user_aor(file_name):
    user_aor_dict = {}
    for lines in open(file_name):
        lines = lines.strip().split('\t')
        user_aor_dict[int(lines[0])] = [int(a) for a in lines[1:]]
    return user_aor_dict

# 读入aor中包含的spu
def read_aor_spu(file_name):
    aor_spu_dict = {}
    for lines in open(file_name):
        lines = lines.strip().split('\t')
        aor_spu_dict[int(lines[0])] = [int(a) for a in lines[1:]]
    return aor_spu_dict

# 读入召回层的物品
def read_recall_items(file_name):
    if 'recall_items_all' in file_name:  # 全部物品参与评估
        return list(range(195244))  # 195244为全部spu的数量
    else:
        recall_items = []
        lines = open(file_name, 'r').readlines()
        for l in lines:
            recall_items.append(int(l.strip()))
        return recall_items


def read_cf(file_name):
    inter_mat = list()
    lines = open(file_name, "r").readlines()
    for l in lines:
        tmps = l.strip()
        inters = [int(i) for i in tmps.split(" ")]

        u_id, pos_ids = inters[0], inters[1:]
        if args.unique_cf:
            pos_ids = list(set(pos_ids))
        else:
            pass
        for i_id in pos_ids:
            inter_mat.append([u_id, i_id])

    return np.array(inter_mat)


def read_user_rebuy(file_name):
    user_rebuy = defaultdict(list)
    lines = open(file_name, "r").readlines()
    for l in lines:
        tmps = l.strip()
        inters = [int(i) for i in tmps.split("\t")]

        u_id, pos_ids = inters[0], inters[1:]
        for i_id in pos_ids:
            user_rebuy[u_id].append(i_id)

    return user_rebuy


def remap_item(train_data, test_data):
    global n_users, n_items
    if args.input_num:
        n_users = args.n_user
        n_items = args.n_item
    else:
        n_users = max(max(train_data[:, 0]), max(test_data[:, 0])) + 1
        n_items = max(max(train_data[:, 1]), max(test_data[:, 1])) + 1

    for u_id, i_id in train_data:
        train_user_set[int(u_id)].append(int(i_id))
    for u_id, i_id in test_data:
        test_user_set[int(u_id)].append(int(i_id))


def read_triplets(file_name):
    global n_entities, n_relations, n_nodes

    can_triplets_np = np.loadtxt(file_name, dtype=np.int32)
    can_triplets_np = np.unique(can_triplets_np, axis=0)

    if args.inverse_r:
        # get triplets with inverse direction like <entity, is-aspect-of, item>
        inv_triplets_np = can_triplets_np.copy()
        inv_triplets_np[:, 0] = can_triplets_np[:, 2]
        inv_triplets_np[:, 2] = can_triplets_np[:, 0]
        inv_triplets_np[:, 1] = can_triplets_np[:, 1] + max(can_triplets_np[:, 1]) + 1
        # consider two additional relations --- 'interact' and 'be interacted'
        can_triplets_np[:, 1] = can_triplets_np[:, 1] + 1
        inv_triplets_np[:, 1] = inv_triplets_np[:, 1] + 1
        # get full version of knowledge graph
        triplets = np.concatenate((can_triplets_np, inv_triplets_np), axis=0)
    else:
        # consider two additional relations --- 'interact'.
        can_triplets_np[:, 1] = can_triplets_np[:, 1] + 1
        triplets = can_triplets_np.copy()

    if args.input_num:
        n_entities = args.n_entity
        n_nodes = args.n_node
        if args.inverse_r:
            n_relations = (args.n_relation + 1) * 2
        else:
            n_relations = args.n_relation + 1
    else:
        n_entities = max(max(triplets[:, 0]), max(triplets[:, 2])) + 1  # including items + users
        n_nodes = n_entities + n_users
        n_relations = max(triplets[:, 1]) + 1

    return triplets


def build_graph(train_data, triplets):
    ckg_graph = nx.MultiDiGraph()
    rd = defaultdict(list)

    logging.info("Begin to load interaction triples ...")
    for u_id, i_id in tqdm(train_data, ascii=True):
        rd[0].append([u_id, i_id])

    logging.info("\nBegin to load knowledge graph triples ...")
    for h_id, r_id, t_id in tqdm(triplets, ascii=True):
        ckg_graph.add_edge(h_id, t_id, key=r_id)
        rd[r_id].append([h_id, t_id])

    return ckg_graph, rd


def build_sparse_relational_graph(relation_dict):
    def _bi_norm_lap(adj):
        # D^{-1/2}AD^{-1/2}
        rowsum = np.array(adj.sum(1))

        d_inv_sqrt = np.power(rowsum, -0.5).flatten()
        d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.
        d_mat_inv_sqrt = sp.diags(d_inv_sqrt)

        # bi_lap = adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt)
        bi_lap = d_mat_inv_sqrt.dot(adj).dot(d_mat_inv_sqrt)
        return bi_lap.tocoo()

    def _si_norm_lap(adj):
        # D^{-1}A
        rowsum = np.array(adj.sum(1))

        d_inv = np.power(rowsum, -1).flatten()
        d_inv[np.isinf(d_inv)] = 0.
        d_mat_inv = sp.diags(d_inv)

        norm_adj = d_mat_inv.dot(adj)
        return norm_adj.tocoo()

    adj_mat_list = []
    logging.info("Begin to build sparse relation matrix ...")
    for r_id in tqdm(relation_dict.keys()):
        np_mat = np.array(relation_dict[r_id])
        if r_id == 0:
            cf = np_mat.copy()
            cf[:, 1] = cf[:, 1] + n_users  # [0, n_items) -> [n_users, n_users+n_items)
            vals = [1.] * len(cf)
            adj = sp.coo_matrix((vals, (cf[:, 0], cf[:, 1])), shape=(n_nodes, n_nodes))
        else:
            vals = [1.] * len(np_mat)
            adj = sp.coo_matrix((vals, (np_mat[:, 0], np_mat[:, 1])), shape=(n_nodes, n_nodes))
        adj_mat_list.append(adj)

    norm_mat_list = [_bi_norm_lap(mat) for mat in adj_mat_list]
    mean_mat_list = [_si_norm_lap(mat) for mat in adj_mat_list]
    # interaction: user->item, [n_users, n_entities]
    norm_mat_list[0] = norm_mat_list[0].tocsr()[:n_users, n_users:].tocoo()
    mean_mat_list[0] = mean_mat_list[0].tocsr()[:n_users, n_users:].tocoo()

    return adj_mat_list, norm_mat_list, mean_mat_list


def load_data(model_args):
    global args
    args = model_args
    directory = args.data_path + args.dataset + '/'

    logging.info('reading train and test user-item set ...')
    train_cf = read_cf(directory + args.train_set+'.txt')
    test_cf = read_cf(directory + args.test_set+'.txt')
    remap_item(train_cf, test_cf)

    user_rebuy = {}

    recall_items = read_recall_items(directory + args.recall_set+'.txt')

    user_aor_dict = {}
    aor_spu_dict = {}

    logging.info('combinating train_cf and kg data ...')
    triplets = read_triplets(directory + args.kg + '.txt')

    logging.info('building the graph ...')
    graph, relation_dict = build_graph(train_cf, triplets)

    logging.info('building the adj mat ...')
    adj_mat_list, norm_mat_list, mean_mat_list = build_sparse_relational_graph(relation_dict)

    n_params = {
        'n_users': int(n_users),
        'n_items': int(n_items),
        'n_entities': int(n_entities),
        'n_nodes': int(n_nodes),
        'n_relations': int(n_relations)
    }
    user_dict = {
        'train_user_set': train_user_set,
        'test_user_set': test_user_set
    }

    return train_cf, test_cf, user_dict, n_params, graph, \
           [adj_mat_list, norm_mat_list, mean_mat_list], recall_items, user_rebuy, user_aor_dict, aor_spu_dict

